import asyncio
import socket
import struct
import time
import random
import string
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from typing import Dict, Optional, Set
import threading

try:
	import uvloop
	asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
	print("🚀 使用 uvloop 加速异步事件循环（仅 Linux/macOS 有效）")
except ImportError:
	print("⚠️ 未安装 uvloop，可使用 'pip install uvloop' 提升性能（仅支持 Linux/macOS）")
except RuntimeError as e:
	if "Windows" in str(e):
		print("⚠️ uvloop 不支持 Windows，使用原生 asyncio 事件循环")
	else:
		print(f"⚠️ uvloop 初始化失败: {e}")

# ==================== 配置参数 ====================
DID_PREFIX = "BOTDDD"  # 设备 DID 前缀

TOTAL_DID_COUNT = 60000  # 总共模拟的设备数量
BATCH_CONNECT_SIZE = 1000  # 每批连接的设备数
BATCH_CONNECT_INTERVAL = 0.2  # 批次连接间隔（秒）

TCP_HOST = "8.153.200.98"  # TCP 服务器地址
TCP_PORT = 50000  # TCP 端口
TCP_HEARTBEAT_INTERVAL = 180  # 心跳包发送间隔（秒）

UDP_HOST = TCP_HOST  # UDP 唤醒服务器地址
UDP_PORT = 50001  # UDP 端口
UDP_BATCH_SIZE = TOTAL_DID_COUNT//10  # 每轮唤醒的设备数量，按10% 比例
WAKEUP_INTERVAL = 2  # 唤醒间隔（秒）
UDP_TIMEOUT = 5  # UDP 发送超时时间
UDP_THREADS = 500  # UDP 发送线程数

ENCRYPTION_KEY = "RSBOTDWakeupKey@"  # 加密密钥
ENCRYPTION_TYPE = 256  # 加密类型

RUN_MINUTES = 10 * 24 * 60  # 运行时长（分钟），例如 10 小时

# ==================== 全局变量 ====================
# 存储正在等待响应的 DID 及其发送时间（UDP 发送时间）
pending_wakeup: Dict[str, float] = {}

# 存储已收到响应的设备延迟结果
latency_results: Dict[str, Dict] = {}

# 当前所有活跃的 TCP 客户端
active_clients: Set['TCPDeviceClient'] = set()

# 已成功建立 TCP 连接的设备计数
connected_count = 0
total_connected = 0

# 上一批次唤醒的设备列表（用于统计响应）
last_batch_dids: Set[str] = set()

# 用于线程安全操作的锁
pending_lock = threading.Lock()  # 保护 pending_wakeup 和 latency_results
active_clients_lock = threading.Lock()  # 保护 active_clients
connect_lock = threading.Lock()  # 保护 connected_count
last_batch_lock = threading.Lock()  # 保护 last_batch_dids

try:
	from string_enc import iPN_StringEnc
except ImportError:
	print("❌ 错误：请确保 string_enc.py 存在并包含 iPN_StringEnc 函数")
	exit(1)


# ==================== 工具函数 ====================
def generate_did_list() -> list:
	"""
	生成指定数量的设备唯一标识（DID）
	格式：IOTGDD-000001-ABCDEF
	"""
	chars = string.ascii_uppercase
	suffix = ''.join(random.choices(chars, k=5))  # 随机5位后缀
	print(f"  生成设备后缀: {suffix}")
	return [f"{DID_PREFIX}-{i:06d}-{suffix}" for i in range(TOTAL_DID_COUNT)]


def create_udp_packet(did: str) -> bytes:
	"""
	构建 UDP 唤醒包
	包格式：4字节魔数 + 4字节消息类型 + 4字节长度 + DID 数据
	"""
	payload = did.encode('utf-8')
	msg_len = len(payload)
	header = struct.pack(">III", 0xFFFFFFA1, 3, msg_len)  # 大端格式
	return header + payload


def send_udp_wakeup(did: str) -> bool:
	"""
	发送单个 UDP 唤醒包
	返回是否发送成功
	"""
	packet = create_udp_packet(did)
	try:
		with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
			sock.settimeout(UDP_TIMEOUT)
			send_time = time.time()
			sock.sendto(packet, (UDP_HOST, UDP_PORT))
			# 记录发送时间，等待响应
			with pending_lock:
				pending_wakeup[did] = send_time
		return True
	except Exception as e:
		print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ❌ UDP 发送失败 | DID: {did} | 错误: {e}")
		return False


# ==================== TCP 客户端类 ====================
class TCPDeviceClient:
	def __init__(self, did: str):
		self.did = did
		_, self.encrypted_data = iPN_StringEnc(ENCRYPTION_KEY, did, ENCRYPTION_TYPE)
		if isinstance(self.encrypted_data, str):
			self.encrypted_data = self.encrypted_data.encode('utf-8')
		self.reader: Optional[asyncio.StreamReader] = None
		self.writer: Optional[asyncio.StreamWriter] = None
		self.connected = False  # 连接状态标志

	async def connect(self):
		"""
		异步连接 TCP 服务器
		成功后加入活跃客户端集合，并启动心跳和监听任务
		"""
		try:
			self.reader, self.writer = await asyncio.wait_for(
				asyncio.open_connection(TCP_HOST, TCP_PORT),
				timeout=10
			)
			self.connected = True
			# 加入全局活跃客户端集合
			with active_clients_lock:
				active_clients.add(self)
			# 更新连接计数
			with connect_lock:
				global connected_count, total_connected
				connected_count += 1
				total_connected += 1
			if connected_count % 1000 == 0 or connected_count == TOTAL_DID_COUNT:
				now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
				print(f"[{now}] ✅ 已连接 {connected_count}/{TOTAL_DID_COUNT}")
			# 启动心跳和监听任务
			asyncio.create_task(self._start_heartbeat())
			asyncio.create_task(self._listen())
		except Exception as e:
			print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ❌ 连接失败 | DID: {self.did} | 错误: {e!r}")

	async def _start_heartbeat(self):
		"""
		错峰发送心跳包
		避免所有设备同时发送导致网络拥塞
		"""
		initial_delay = random.uniform(0, TCP_HEARTBEAT_INTERVAL)
		await asyncio.sleep(initial_delay)
		while self.connected:
			if self.writer and not self.writer.is_closing():
				try:
					self.writer.write(self.encrypted_data)
					await self.writer.drain()  # 确保数据发送
				except Exception as e:
					error_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
					print(f"[{error_time}] ❌ 心跳发送失败 | DID: {self.did} | 错误: {e}")
					self.connected = False
					break
			await asyncio.sleep(TCP_HEARTBEAT_INTERVAL)

	async def _listen(self):
		"""
		持续监听 TCP 服务端响应
		收到数据后计算延迟，并记录结果
		"""
		try:
			while self.connected:
				data = await self.reader.read(1024)
				if not data:
					break
				recv_time = time.time()
				# 检查该设备是否在等待响应中
				with pending_lock:
					if self.did in pending_wakeup:
						send_time = pending_wakeup[self.did]
						delay = recv_time - send_time
						# 保存延迟结果
						latency_results[self.did] = {
							"DID": self.did,
							"UDP_Send_Time": datetime.fromtimestamp(send_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
							"TCP_Recv_Time": datetime.fromtimestamp(recv_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
							"Latency_ms": round(delay * 1000, 2)
						}
						# 从待响应列表中移除
						del pending_wakeup[self.did]
					else:
						print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ⚠️ 收到未请求响应 | DID: {self.did}")
		except Exception as e:
			print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ❌ TCP 监听异常 | DID: {self.did} | 错误: {e}")
		finally:
			self.connected = False
			with active_clients_lock:
				active_clients.discard(self)
			if self.writer and not self.writer.is_closing():
				self.writer.close()
				try:
					await self.writer.wait_closed()
				except:
					pass


# ==================== 唤醒与统计函数 ====================
def trigger_udp_wakeup_batch():
	"""
	发送一批 UDP 唤醒任务
	从当前活跃设备中随机选择一批进行唤醒
	"""
	# 获取当前所有已连接的设备
	with active_clients_lock:
		live_dids = [client.did for client in active_clients if client.connected]

	if not live_dids:
		print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] ⚠️  当前无活跃设备，跳过唤醒")
		return

	# 随机选择一批设备
	selected_dids = random.sample(live_dids, k=min(UDP_BATCH_SIZE, len(live_dids)))
	submit_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

	print(f"[{submit_time_str}] ⏱️  提交 UDP 唤醒任务: 数量={len(selected_dids)}")

	# 记录本批次唤醒的设备（用于后续统计）
	with last_batch_lock:
		global last_batch_dids
		last_batch_dids = set(selected_dids)

	# 使用线程池发送 UDP 包
	with ThreadPoolExecutor(max_workers=UDP_THREADS) as executor:
		futures = [executor.submit(send_udp_wakeup, did) for did in selected_dids]
		send_success = sum(1 for f in as_completed(futures) if f.result())

	done_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
	print(f"[{done_time_str}] 🚀 UDP 唤醒完成 | 发送: {len(selected_dids)} | 成功发送: {send_success}")


def print_batch_statistics():
	"""
	统计并打印上一批唤醒任务的响应情况
	包括：响应数、延迟统计、QPS 等
	并打印未响应的 DID 列表
	"""
	with last_batch_lock:
		if not last_batch_dids:
			return
		selected_dids = last_batch_dids.copy()
	# print(F"last_batch_dids个数：{len(last_batch_dids)}")
	responded_dids = []
	latencies = []

	# 统计哪些设备已响应
	with pending_lock:
		for did in selected_dids:
			if did in latency_results:
				latencies.append(latency_results[did]["Latency_ms"])
				responded_dids.append(did)

	# 计算未响应的设备
	unresponded_dids = [did for did in selected_dids if did not in responded_dids]
	responded_count = len(responded_dids)
	total_count = len(selected_dids)

	# 打印统计信息
	print("📈 本批次统计:")
	print(f"   发送: {total_count} | 成功响应: {responded_count} | 未响应: {len(unresponded_dids)}")

	if latencies:
		# 计算统计值
		avg = sum(latencies) / len(latencies)
		min_lat = min(latencies)
		max_lat = max(latencies)
		sorted_lat = sorted(latencies)
		p95_idx = int(0.95 * len(sorted_lat))
		p95 = sorted_lat[p95_idx] if p95_idx < len(sorted_lat) else sorted_lat[-1]
		p99_idx = int(0.99 * len(sorted_lat))
		p99 = sorted_lat[p99_idx] if p99_idx < len(sorted_lat) else sorted_lat[-1]
		qps = responded_count / WAKEUP_INTERVAL

		print(f"   平均={avg:.2f}ms | 最小={min_lat:.2f}ms | 最大={max_lat:.2f}ms")
		print(f"   P95={p95:.2f}ms | P99={p99:.2f}ms | QPS={qps:.2f}")
	else:
		print("   ❌ 无响应数据")

	# ✅ 打印未响应的 DID（可选：只在未响应数 > 0 时打印）
	if unresponded_dids:
		print(f"   📛 未响应 DID 列表 ({len(unresponded_dids)} 个):")
		# 为了避免输出太长，可以只打印前 10 个，或全部打印
		# for did in unresponded_dids:
		# 	print(f"     🔴 {did}")
		# 如果设备太多，也可以只打印前几项 + 省略号
		for did in unresponded_dids[:10]:
			print(f"     🔴 {did}")
		if len(unresponded_dids) > 10:
			print(f"     ... 还有 {len(unresponded_dids) - 10} 个")


# ==================== 主函数 ====================
async def main():
	print(f"🚀 物联网设备唤醒延迟测试平台启动")
	print(f"  TCP 长连接数: {TOTAL_DID_COUNT}")
	print(f"  分批连接: {BATCH_CONNECT_SIZE} 个/批")
	print(f"  UDP 唤醒: 每 {WAKEUP_INTERVAL} 秒唤醒 {UDP_BATCH_SIZE} 个设备")
	print(f"  运行时长: {RUN_MINUTES} 分钟")

	# 生成设备 DID 列表
	did_list = generate_did_list()
	tcp_clients = [TCPDeviceClient(did) for did in did_list]

	print("🟡 正在分批建立 TCP 连接...")
	# 分批连接，避免瞬时压力过大
	for i in range(0, len(tcp_clients), BATCH_CONNECT_SIZE):
		batch = tcp_clients[i:i + BATCH_CONNECT_SIZE]
		for client in batch:
			asyncio.create_task(client.connect())
		await asyncio.sleep(BATCH_CONNECT_INTERVAL)

	print(f"✅ 已提交 {len(tcp_clients)} 个连接任务，正在后台连接...")

	start_time = time.time() + TCP_HEARTBEAT_INTERVAL
	end_time = start_time + RUN_MINUTES * 60
	next_wakeup_time = start_time

	try:
		while time.time() < end_time:
			sleep_time = next_wakeup_time - time.time()
			if sleep_time > 0:
				await asyncio.sleep(sleep_time)

			# ✅ 先统计上一批响应情况
			if 'last_wakeup_time' in locals():
				print_batch_statistics()

			# ✅ 再发送新一批唤醒
			trigger_udp_wakeup_batch()
			last_wakeup_time = time.time()

			next_wakeup_time += WAKEUP_INTERVAL

	except KeyboardInterrupt:
		print("\n🛑 用户中断，正在打印最终统计...")
		print_batch_statistics()
	finally:
		print("\n✅ 测试结束，最终统计:")
		print_batch_statistics()


if __name__ == "__main__":
	asyncio.run(main())